/* OTA example

   This example code is in the Public Domain (or CC0 licensed, at your option.)

   Unless required by applicable law or agreed to in writing, this
   software is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
   CONDITIONS OF ANY KIND, either express or implied.
*/
#include <string.h>
#include "freertos/FreeRTOS.h"
#include "freertos/task.h"
#include "esp_system.h"
#include "esp_event.h"
#include "esp_log.h"
#include "esp_ota_ops.h"
#include "esp_flash_partitions.h"
#include "esp_partition.h"
#include "nvs.h"
#include "nvs_flash.h"
#include "driver/gpio.h"
#include "protocol_examples_common.h"
#include "errno.h"
#include "esp_wifi.h"
#include "ota.h"

#define BUFFSIZE 1024
#define HASH_LEN 32 /* SHA-256 digest length */

#define OTA_SECH_IDLE 0
#define OTA_REQUEST_PACKET 1
#define OTA_DOWNLOAD_PACKET 2
#define OTA_PROGRAM 3
#define OTA_SECH_COMPLETE 4

typedef struct
{
    char version[16];
    uint16_t total_packet_num;
    uint32_t image_size;
    uint8_t image_sha256[32];
} ota_image_info_t;

static const char *TAG = "ota";
/*an ota data write buffer ready to write to the flash*/
static char ota_write_data[BUFFSIZE + 1] = {0};
static uint16_t packet_num;
static uint8_t ota_scheduler;
static ota_image_info_t image_info;
static uint32_t timer;
static void ota_request_packet(void)
{
    uint8_t packet_num_c[2];
    msg_t msg = {0};
    LOGI(TAG, "Request packet: %d\n", packet_num);
    packet_num_c[0] = packet_num;
    packet_num_c[1] = packet_num >> 8;
    msg.len = protocal_packet(runtime_profile.device_id, 0x09, msg.data, 2, packet_num_c);
    queue_add(&msg, true);
    ota_scheduler = OTA_DOWNLOAD_PACKET;
}

void ota_request(uint8_t *pData, uint16_t len)
{
    esp_app_desc_t running_app_info;
    if (ota_scheduler != OTA_SECH_IDLE)
    {
        ESP_LOGW(TAG, "OTA not idle: %d\n", ota_scheduler);
        //TODO: notify server
        return;
    }
    const esp_partition_t *running = esp_ota_get_running_partition();
    if (esp_ota_get_partition_description(running, &running_app_info) == ESP_OK)
    {
        if (0 != strncpy(running_app_info.version, pData, strlen(running_app_info.version)))
        {
            ESP_LOGW(TAG, "target version same with current version:%s\n", running_app_info.version);
            //TODO: notify server
            return;
        }
    }
    memset(&image_info, 0, sizeof(image_info));
    memcpy(image_info.version, pData, 13);
    image_info.total_packet_num = pData[13] | pData[14] << 8;
    if (image_info.total_packet_num == 0 || image_info.total_packet_num > 2048)
    {
        ESP_LOGW(TAG, "total packet num error: %d\n", image_info.total_packet_num);
        //TODO: notify server
        return;
    }
    image_info.image_size = pData[15] | pData[16] << 8 | pData[17] << 16 | pData[18] << 24;
    if (image_info.image_size >= image_info.total_packet_num * 1024)
    {
        ESP_LOGW(TAG, "total image size error: %d\n", image_info.image_size);
        //TODO: notify server
        return;
    }
    memcpy(image_info.image_sha256, pData[19], 32);

    esp_err_t err = nvs_flash_init();
    if (err == ESP_ERR_NVS_NO_FREE_PAGES || err == ESP_ERR_NVS_NEW_VERSION_FOUND)
    {
        // OTA app partition table has a smaller NVS partition size than the non-OTA
        // partition table. This size mismatch may cause NVS initialization to fail.
        // If this happens, we erase NVS partition and initialize NVS again.
        ESP_ERROR_CHECK(nvs_flash_erase());
        err = nvs_flash_init();
    }
    ESP_ERROR_CHECK(err);
    ota_scheduler = OTA_SECH_ONGOING;
    ota_request_packet();
    //ESP_ERROR_CHECK(esp_netif_init());
    //ESP_ERROR_CHECK(esp_event_loop_create_default());
}

void ota_packet_handler(uint8_t *pData, uint16_t len)
{

    esp_err_t err;
    /* update handle : set by esp_ota_begin(), must be freed via esp_ota_end() */
    esp_ota_handle_t update_handle = 0;
    const esp_partition_t *update_partition = NULL;

    const esp_partition_t *configured = esp_ota_get_boot_partition();
    const esp_partition_t *running = esp_ota_get_running_partition();

    if (configured != running)
    {
        ESP_LOGW(TAG, "Configured OTA boot partition at offset 0x%08x, but running from offset 0x%08x",
                 configured->address, running->address);
        ESP_LOGW(TAG, "(This can happen if either the OTA boot data or preferred boot image become corrupted somehow.)");
    }
    ESP_LOGI(TAG, "Running partition type %d subtype %d (offset 0x%08x)",
             running->type, running->subtype, running->address);

    update_partition = esp_ota_get_next_update_partition(NULL);
    ESP_LOGI(TAG, "Writing to partition subtype %d at offset 0x%x", update_partition->subtype, update_partition->address);
    assert(update_partition != NULL);

    int binary_file_length = 0;
    /*deal with all receive packet*/
    bool image_header_was_checked = false;
    while (1)
    {
        //int data_read = esp_http_client_read(client, ota_write_data, BUFFSIZE);
        if (data_read < 0)
        {
        }
        else if (data_read > 0)
        {
            if (image_header_was_checked == false)
            {
                esp_app_desc_t new_app_info;
                if (data_read > sizeof(esp_image_header_t) + sizeof(esp_image_segment_header_t) + sizeof(esp_app_desc_t))
                {
                    // check current version with downloading
                    memcpy(&new_app_info, &ota_write_data[sizeof(esp_image_header_t) + sizeof(esp_image_segment_header_t)], sizeof(esp_app_desc_t));
                    ESP_LOGI(TAG, "New firmware version: %s", new_app_info.version);

                    esp_app_desc_t running_app_info;
                    if (esp_ota_get_partition_description(running, &running_app_info) == ESP_OK)
                    {
                        ESP_LOGI(TAG, "Running firmware version: %s", running_app_info.version);
                    }

                    const esp_partition_t *last_invalid_app = esp_ota_get_last_invalid_partition();
                    esp_app_desc_t invalid_app_info;
                    if (esp_ota_get_partition_description(last_invalid_app, &invalid_app_info) == ESP_OK)
                    {
                        ESP_LOGI(TAG, "Last invalid firmware version: %s", invalid_app_info.version);
                    }

                    // check current version with last invalid partition
                    if (last_invalid_app != NULL)
                    {
                        if (memcmp(invalid_app_info.version, new_app_info.version, sizeof(new_app_info.version)) == 0)
                        {
                            ESP_LOGW(TAG, "New version is the same as invalid version.");
                            ESP_LOGW(TAG, "Previously, there was an attempt to launch the firmware with %s version, but it failed.", invalid_app_info.version);
                            ESP_LOGW(TAG, "The firmware has been rolled back to the previous version.");
                            http_cleanup(client);
                            infinite_loop();
                        }
                    }
                    if (memcmp(new_app_info.version, running_app_info.version, sizeof(new_app_info.version)) == 0)
                    {
                        ESP_LOGW(TAG, "Current running version is the same as a new. We will not continue the update.");
                        //TODO: donot upgrade
                    }
                    else
                    {
                        image_header_was_checked = true;

                        err = esp_ota_begin(update_partition, OTA_SIZE_UNKNOWN, &update_handle);
                        if (err != ESP_OK)
                        {
                            ESP_LOGE(TAG, "esp_ota_begin failed (%s)", esp_err_to_name(err));
                        }
                        else
                        {
                            ESP_LOGI(TAG, "esp_ota_begin succeeded");
                        }
                    }
                }
                else
                {
                    ESP_LOGE(TAG, "received package is not fit len");
                }
            }
            err = esp_ota_write(update_handle, (const void *)ota_write_data, data_read);
            if (err != ESP_OK)
            {
            }
            binary_file_length += data_read;
            ESP_LOGD(TAG, "Written image length %d", binary_file_length);
        }
        else if (data_read == 0)
        {
            /*
            * As esp_http_client_read never returns negative error code, we rely on
            * `errno` to check for underlying transport connectivity closure if any
            */
            if (errno == ECONNRESET || errno == ENOTCONN)
            {
                ESP_LOGE(TAG, "Connection closed, errno = %d", errno);
                break;
            }
            if (esp_http_client_is_complete_data_received(client) == true)
            {
                ESP_LOGI(TAG, "Connection closed");
                break;
            }
        }
    }
    ESP_LOGI(TAG, "Total Write binary data length: %d", binary_file_length);

    err = esp_ota_end(update_handle);
    if (err != ESP_OK)
    {
        if (err == ESP_ERR_OTA_VALIDATE_FAILED)
        {
            ESP_LOGE(TAG, "Image validation failed, image is corrupted");
        }
        ESP_LOGE(TAG, "esp_ota_end failed (%s)!", esp_err_to_name(err));
    }

    err = esp_ota_set_boot_partition(update_partition);
    if (err != ESP_OK)
    {
        ESP_LOGE(TAG, "esp_ota_set_boot_partition failed (%s)!", esp_err_to_name(err));
    }
    ESP_LOGI(TAG, "Prepare to restart system!");
    esp_restart();
    return;
}

static bool diagnostic(void)
{
    return true;
}

void ota_task(void)
{
    switch (ota_scheduler)
    {
    case OTA_SECH_IDLE:
        break;
    case OTA_REQUEST_PACKET:
        ota_request_packet();
        break;
    case OTA_DOWNLOAD_PACKET:
        break;
    case OTA_PROGRAM:
        break;
    case OTA_SECH_COMPLETE:
        break;
    }
}